from sklearn.model_selection import LeaveOneOut, StratifiedKFold
from sklearn.base import clone
from classifier.dataset import DataPreprocessor, DataLoader
from classifier.model import SVMModel, DecisionTreeModel, NeuralNetworkModel, RandomForest
from classifier.visual import ROCVisualizer


def main():
    data = DataLoader(shuffle=True)
    data.load_csv("./both.csv")

    preprocessor = DataPreprocessor()
    models = {
        "SVM": SVMModel(preprocessor),
        "DecisionTree": DecisionTreeModel(preprocessor),
        "NeuralNetwork": NeuralNetworkModel(preprocessor),
        "RandomForest": RandomForest(preprocessor)
    }
    results = {}
    for name, model in models.items():
        # 留一法训练模型
        model.train(data.x, data.y)
        # 5折交叉验证评估模型
        cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
        for i, (train_idx, test_idx) in enumerate(cv.split(data.x, data.y)):
            model.best_model.fit(data.x[train_idx], data.y[train_idx])
            model.evaluate(data.x[test_idx], data.y[test_idx])
        # 收集评估结果
        results[name] = model.get_metrics()

        # 打印结果
        print(f"=============={name}============")
        print(f"Best Parameters: {model.best_params_}")
        for metric, value in results[name].items():
            if metric not in ["fpr", "tpr"]:
                print(f"{metric.capitalize()}: {value:.4f}")

    # 可视化
    ROCVisualizer.plot_auc_curves(results)


if __name__ == "__main__":
    main()
